#include <ctime>
#include <cstdlib>
#include <unistd.h>
#include <iostream>
#include <iomanip>

#include "Input/PicParams.h"
#include "Parallel/ParallelMPIFactory.h"
#include "Grid/GridFactory.h"
#include "Species/SpeciesFactory.h"
#include "ParticleSource/ParticleSourceFactory.h"
#include "Collision/CollisionFactory.h"
#include "PhysicalField/PhysicalFieldFactory.h"
#include "FieldSolver/SolverFactory.h"
#include "IO/IOFactory.h"
#include "PSI/PSIFactory.h"
#include "Diagnostic/DiagnosticFactory.h"
#include "Tool/Timer.h"
#include "Tool/Macros.h"
#include "Tool/Log.h"

#ifdef  __PICDEBUG
#include <gperftools/profiler.h>
#endif

using namespace std;


int main (int argc, char* argv[])
{
    //pmpi_base is used to init log
    ParallelMPI* pmpi_base = new ParallelMPI(&argc, &argv);

    Log* log = new Log(pmpi_base);

    log_title("Creating PicParams");
    PicParams* params = new PicParams();
    pmpi_base->barrier();

    log_title("Creating ParallelMPI");
    ParallelMPI* pmpi = ParallelMPIFactory::create(params, pmpi_base);
    pmpi->barrier();

    log_title("Creating IO");
    IO* io = IOFactory::create(params, pmpi);
    pmpi->barrier();
    
    log_title("Creating Grid");
    Grid* grid = GridFactory::create(params, pmpi, io);

    log_title("Creating PhysicalField");
    PhysicalField* fields = PhysicalFieldFactory::create(params, pmpi);

    log_title("Creating Species");
    vector<Species*> vecSpecies = SpeciesFactory::createVector(params, pmpi);

    log_title("Creating Collision");
    vector<Collision*> vecCollision = CollisionFactory::create(params, pmpi, &vecSpecies);

    log_title("Creating PSI");
    vector<PSI*> vecPSI = PSIFactory::create(params, pmpi, &vecSpecies);

    log_title("Creating Diagnostic");
    Diagnostic*  diag  = DiagnosticFactory::create(params, pmpi, grid, fields, &vecSpecies, &vecCollision, &vecPSI);

    log_title("Creating ParticleSource");
    vector<ParticleSource*> vecParticleSource = ParticleSourceFactory::create(params, pmpi, fields, &vecSpecies);

    log_title("Creating Solver");
    Solver* solver =  SolverFactory::create(params, pmpi, grid, fields, diag);

    log_title("Init IO environment");
    io->init(grid, fields, &vecSpecies, &vecParticleSource, &vecCollision, &vecPSI, diag);
    pmpi->barrier();

    //Initialize simulation time
    int step_start = io->step_start;
    int step_stop  = params->timestep_simulation;

    int itime = step_start;


    //creating timer
    vector<Timer> timer(15);
    timer[0].init(pmpi, "Total time");
    timer[1].init(pmpi, "EmitLoad");
    timer[2].init(pmpi, "Collide");
    timer[3].init(pmpi, "Move");
    timer[4].init(pmpi, "MPI Exchange Particle");
    timer[5].init(pmpi, "Absorb paritcle (2D)");
    timer[6].init(pmpi, "Project Particle");
    timer[7].init(pmpi, "PSI");
    timer[8].init(pmpi, "Diagnostic");
    timer[9].init(pmpi, "Fields Solve");
    timer[10].init(pmpi,"Write IO");
    timer[11].init(pmpi,"Field solver init");
    timer[12].init(pmpi,"Particle sort");
    //timer[13].init(pmpi,"Correct particles");
    timer[14].init(pmpi,"Fields Sovle petsc");

    timer[0].restart();

    log_title("Project Particles first time before PIC loop");
    for(int ispec = 0; ispec < params->species_param.size(); ispec++)
    {
        fields->restartRhoJs(ispec);
        vecSpecies[ispec]->project(fields);
    }
    fields->computeTotalRhoJ();

    log_title("Solve the field first time before PIC loop");
    if(params->is_self_consistent_electric_field)
    {
        (*solver)(pmpi, fields, diag);
    }
    else
    {
        io->read(params, pmpi, fields, vecSpecies, diag, itime);
    }
    pmpi->barrier();

    log_title("Output the field first time before PIC loop");
    fields->incrementAvgFields(itime);
    fields->gatherFields(pmpi);
    fields->gatherAvgFields(pmpi);
    diag->run(pmpi, grid, vecSpecies, fields, vecPSI, itime);
    io->write(params, pmpi, fields, vecSpecies, diag, itime);

    //google perftools for PICDEBUG
    #ifdef  __PICDEBUG
    long long mpi_rk = pmpi->getRank();
    string prof_name = "./liz.prof_" + to_string( mpi_rk );
    ProfilerStart(prof_name.c_str());
    #endif


    //============================= start PIC main loop ===========================================
    log_title("Time-Loop is started: number of output = " << params->timestep_simulation / params->timestep_io);

    if(params->method == "explicit")
    {
        while(itime <= step_stop)
        {
            itime++;
            log_debug<<"timestep = "<<itime;

            //================== EmitLoad =========================================
            timer[1].restart();
            log_debug<<"Particle Source";
            for(int i_ps = 0 ; i_ps < vecParticleSource.size(); i_ps++)
            {
                vecParticleSource[i_ps]->emitLoad(itime);
            }
            timer[1].update();

            //================== Collide =========================================
            timer[2].restart();
            log_debug<<"Collide";
            for(int icoll = 0 ; icoll < vecCollision.size(); icoll++)
            {
                vecCollision[icoll]->collide(fields, vecSpecies, diag, itime);
            }
            timer[2].update();

            //================== Interpolate and Move ===============================
            timer[3].restart();
            log_debug<<"Move Particle";
            for (int ispec=0 ; ispec<params->species_param.size(); ispec++)
            {
                if(itime % params->species_param[ispec].timestep_frozen != 0) continue;
                vecSpecies[ispec]->move(fields);
            }
            timer[3].update();

            //================== Interpolate and Move ===============================
            timer[12].restart();
            log_debug<<"Move Particle";
            for (int ispec=0 ; ispec<params->species_param.size(); ispec++)
            {
                if(itime % params->species_param[ispec].timestep_frozen != 0) continue;
                vecSpecies[ispec]->sort_particles();
                vecSpecies[ispec]->correct_particles();
            }
            timer[12].update();

            //================== MPI Exchange Particle ============================================
            timer[4].restart();
            log_debug<<"MPI Exchange Particle";
            for(int ispec = 0; ispec < params->species_param.size(); ispec++)
            {
                if(itime % params->species_param[ispec].timestep_frozen != 0) continue;
                for(int i_dim = 0; i_dim < params->n_dim_field; i_dim++)
                {
                    pmpi->exchange_particles(vecSpecies[ispec], params, i_dim);
                }
            }
            pmpi->barrier();
            timer[4].update();

            //================== Project Particle ========================================
            timer[6].restart();
            log_debug<<"Project Particle";
            for (int ispec=0 ; ispec<params->species_param.size(); ispec++)
            {
                if(itime % params->species_param[ispec].timestep_frozen != 0) continue;
                fields->restartRhoJs(ispec);
                vecSpecies[ispec]->project(fields);
            }
            pmpi->barrier();
            timer[6].update();

            //================== Run Diagnostic =============================================
            timer[8].restart();
            log_debug<<"Run Diagnostic";
            diag->run(pmpi, grid, vecSpecies, fields, vecPSI, itime);
            timer[8].update();

            //================== Plasma Surface Interacton ==================================
            timer[7].restart();
            log_debug<<"PSI";
            for (int ipsi=0 ; ipsi<vecPSI.size(); ipsi++)
            {
                vecPSI[ipsi]->perform_psi(itime);
            }
            pmpi->barrier();
            timer[7].update();

            //================== Solve Electromagnetic Fields ===============================
            timer[9].restart();
            log_debug<<"Solve Fields";
            fields->restartRhoJ();
            fields->computeTotalRhoJ();
            fields->gatherFields(pmpi);
            if(params->is_self_consistent_electric_field)
            {
                (*solver)(pmpi, fields, diag);
            }
            timer[9].update();

            //================== Write IO ====================================================
            timer[10].restart();
            log_debug<<"Write IO";
            if(params->timestep_average)
            {
                fields->incrementAvgFields(itime);
            }
            if(itime % params->timestep_io == 0){
                fields->gatherAvgFields(pmpi);
                log_info<<"time step = "<<itime/params->timestep_io;
            }
            io->write(params, pmpi, fields, vecSpecies, diag, itime);
            if(itime % params->timestep_restore == 0)
            {
                io->storeP(params, pmpi, vecSpecies, itime);
            }
            timer[10].update();

        }//END of the time loop
    }
    else if(params->method == "implicit")
    {

    }

    #ifdef  __PICDEBUG
    ProfilerStop();
    #endif


    //======================= output time profile =============================================
    timer[0].update();
    log_title("Program time profile");
    for( int i = 0; i < timer.size(); i++)
    {
        if(i == 0)
        {
            log_info<<std::left<<std::setw(25)<<timer[i].name() << " = " << timer[i].getDateTime();
        }
        else 
        {
            log_info<<std::left<<std::setw(25)<<timer[i].name() << " = " << timer[i].getTime() <<" s";
        }
    }

    log_title("Field solver time profile");
    for( int i = 0; i < solver->timer_solver.size(); i++)
    {
        log_info<<solver->timer_solver[i].name() << " = " << solver->timer_solver[i].getTime() <<" s";
    }

    //===================== clean up simulation ============================================ 
    io->endStoreP(params, pmpi, vecSpecies, itime);


    for (int iPS=0 ; iPS<vecParticleSource.size(); iPS++) delete vecParticleSource[iPS];
    vecParticleSource.clear();

    for(int i=0; i<vecCollision.size(); i++) delete vecCollision[i];
    vecCollision.clear();

    for (int ispec=0 ; ispec<vecSpecies.size(); ispec++) delete vecSpecies[ispec];
    vecSpecies.clear();

    delete fields;
    delete solver;
    delete io;
    delete pmpi;

    log_title("end");
    return 0;

}
